package main

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"math/rand"
	"time"

	_ "github.com/lib/pq"
)

const (
	dbHost     = "localhost"
	dbPort     = 5432
	dbUser     = "postgres"
	dbPassword = "123456"
	dbName     = "performance_test"
	tableName  = "test_table"
)

var (
	letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
	db          *sql.DB
	ctx         = context.Background()
)

func main() {
	var err error
	connStr := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
		dbHost, dbPort, dbUser, dbPassword, dbName)

	db, err = sql.Open("postgres", connStr)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	// 创建测试表
	if err := createTestTable(); err != nil {
		log.Fatal(err)
	}

	for i := 0; i < 30; i++ {
		// 插入测试数据
		if err := insertTestData(10000); err != nil { // 插入10000条测试数据
			log.Fatal(err)
		}
	}

	// 执行查询测试
	if err := runQueryTests(); err != nil {
		log.Fatal(err)
	}
}

// 生成随机字符串
func randString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

// 执行查询测试
func runQueryTests() error {
	// 测试各种查询组合
	queries := []struct {
		name string
		sql  string
	}{
		{
			name: "单个字符索引字段查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_char_1 = 'fixedValue'", tableName),
		},
		{
			name: "单个数字索引字段查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_num_1 = 12345", tableName),
		},
		{
			name: "单个时间索引字段查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_time_1 > NOW() - INTERVAL '1 day'", tableName),
		},
		{
			name: "单个数组索引字段查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE 'abcde' = ANY(index_array_1)", tableName),
		},
		{
			name: "字符+数字索引字段组合查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_char_1 = 'fixedValue' AND index_num_1 = 12345", tableName),
		},
		{
			name: "字符+时间索引字段组合查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_char_1 = 'fixedValue' AND index_time_1 > NOW() - INTERVAL '1 day'", tableName),
		},
		{
			name: "字符+数组索引字段组合查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_char_1 = 'fixedValue' AND 'abcde' = ANY(index_array_1)", tableName),
		},
		{
			name: "多字段组合查询",
			sql:  fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE index_char_1 = 'fixedValue' AND index_num_1 = 12345 AND index_time_1 > NOW() - INTERVAL '1 day' AND 'abcde' = ANY(index_array_1)", tableName),
		},
	}

	for _, q := range queries {
		start := time.Now()
		var count int
		err := db.QueryRowContext(ctx, q.sql).Scan(&count)
		if err != nil {
			return fmt.Errorf("查询 %s 失败: %v", q.name, err)
		}
		elapsed := time.Since(start)
		log.Printf("%s - 返回 %d 条记录, 耗时: %v", q.name, count, elapsed)
	}

	return nil
}
